use crate::{
    ir::{
        attribute::AttrValue,
        node::{Node, NodeId},
        value::{Value, ValueId},
    },
    onnx,
};

use alloc::string::{String, ToString};
use alloc::vec::Vec;

use core::fmt;

pub struct Graph {
    pub values: Vec<Value>,
    pub nodes: Vec<Node>,
    pub inputs: Vec<ValueId>,
    pub outputs: Vec<ValueId>,
    pub name_to_value: hashbrown::HashMap<String, ValueId>,
    pub producers: Vec<Option<NodeId>>,
    pub consumers: Vec<Vec<NodeId>>,
}

impl Graph {
    pub fn from_onnx(model: &onnx::ModelProto) -> Self {
        let g = model.graph.as_ref().expect("no graph");
        let m: hashbrown::HashMap<String, ValueId> = Default::default();

        let mut graph = Graph {
            values: Vec::new(),
            nodes: Vec::new(),
            inputs: Vec::new(),
            outputs: Vec::new(),
            name_to_value: m,
            producers: Vec::new(),
            consumers: Vec::new(),
        };

        for vi in &g.input {
            let vi = vi.clone();
            let id = graph.add_value(&vi.name.unwrap());
            graph.inputs.push(id);
        }

        for vi in &g.output {
            let vi = vi.clone();
            let id = graph.add_value(&vi.name.unwrap());
            graph.outputs.push(id);
        }

        for init in &g.initializer {
            let init = init.clone();
            let id = graph.add_value(&init.name.clone().unwrap());
            if let Some(t) = crate::ir::initializer::from_tensor_proto(&init) {
                graph.values[id as usize].constant = Some(t);
            }
        }

        for n in &g.node {
            let n = n.clone();
            let node_id = graph.nodes.len() as u32;
            let mut node = Node {
                id: node_id,
                domain: n.domain.unwrap_or("ai.onnx".to_string()).clone(),
                op_type: n.op_type.unwrap().clone(),
                inputs: Vec::new(),
                outputs: Vec::new(),
                attrs: Vec::new(),
            };

            for name in &n.input {
                if !name.is_empty() {
                    let vid = graph.add_value(name);
                    node.inputs.push(vid);
                    graph.add_consumer(vid, node_id);
                }
            }
            for name in &n.output {
                if !name.is_empty() {
                    let vid = graph.add_value(name);
                    node.outputs.push(vid);
                    graph.producers_resize(vid, node_id);
                }
            }

            for a in &n.attribute {
                let val = match a.r#type {
                    Some(1) => AttrValue::Float(a.f.unwrap()),
                    Some(2) => AttrValue::Int(a.i.unwrap()),
                    Some(3) => AttrValue::Bytes(a.s.clone().unwrap()),
                    Some(4) => AttrValue::Tensor(
                        crate::ir::initializer::try_from_tensor_proto(&a.clone().t.unwrap())
                            .unwrap(),
                    ),
                    Some(6) => AttrValue::Floats(a.floats.clone()),
                    Some(7) => AttrValue::Ints(a.ints.clone()),
                    _ => continue,
                };
                node.attrs.push((a.name.clone().unwrap(), val));
            }

            graph.nodes.push(node);
        }

        graph
    }

    fn add_value(&mut self, name: &str) -> ValueId {
        if let Some(&id) = self.name_to_value.get(name) {
            return id;
        }
        let id = self.values.len() as ValueId;
        self.name_to_value.insert(name.to_string(), id);
        self.values.push(Value {
            id,
            name: name.to_string(),
            ty: None,
            constant: None,
        });
        self.producers.push(None);
        self.consumers.push(Vec::new());
        id
    }

    fn add_consumer(&mut self, vid: ValueId, nid: NodeId) {
        if let Some(c) = self.consumers.get_mut(vid as usize) {
            c.push(nid);
        }
    }

    fn producers_resize(&mut self, vid: ValueId, nid: NodeId) {
        if let Some(p) = self.producers.get_mut(vid as usize) {
            *p = Some(nid);
        }
    }

    pub fn topo_order(&self) -> alloc::vec::Vec<NodeId> {
        use alloc::vec::Vec;

        let n = self.nodes.len();
        let mut indeg: Vec<u32> = alloc::vec![0; n];
        for (i, node) in self.nodes.iter().enumerate() {
            let mut d = 0u32;
            for &vid in &node.inputs {
                if let Some(p) = self.producers.get(vid as usize).and_then(|x| *x) {
                    if p != node.id {
                        d += 1;
                    }
                }
            }
            indeg[i] = d;
        }

        let mut stack: Vec<NodeId> = Vec::new();
        for (i, &d) in indeg.iter().enumerate() {
            if d == 0 {
                stack.push(i as NodeId);
            }
        }

        let mut order: Vec<NodeId> = Vec::with_capacity(n);
        while let Some(nid) = stack.pop() {
            order.push(nid);
            let node = &self.nodes[nid as usize];
            for &ov in &node.outputs {
                if let Some(cons) = self.consumers.get(ov as usize) {
                    for &cid in cons {
                        if cid == nid {
                            continue;
                        }
                        let idx = cid as usize;
                        debug_assert!(indeg[idx] > 0, "indegree underflow for node {}", cid);
                        indeg[idx] -= 1;
                        if indeg[idx] == 0 {
                            stack.push(cid);
                        }
                    }
                }
            }
        }

        if order.len() != n {
            let mut seen = alloc::vec![false; n];
            for &nid in &order {
                seen[nid as usize] = true;
            }
            for i in 0..n {
                if !seen[i] {
                    order.push(i as NodeId);
                }
            }
        }

        order
    }
}

#[derive(Clone, Copy)]
pub struct PrettyOpts {
    pub show_attrs: bool,
    pub show_constants: bool,
    pub const_max_elems: usize,
}

impl Default for PrettyOpts {
    fn default() -> Self {
        Self {
            show_attrs: true,
            show_constants: false,
            const_max_elems: 8,
        }
    }
}

impl Graph {
    pub fn pretty(&self, opts: PrettyOpts) -> alloc::string::String {
        use alloc::format;

        let mut s = alloc::string::String::new();
        s.push_str("inputs: [");
        for (i, &vid) in self.inputs.iter().enumerate() {
            if i > 0 {
                s.push_str(", ");
            }
            s.push_str(&format!("{}", self.values[vid as usize].name));
        }
        s.push_str("]\n");

        s.push_str("outputs: [");
        for (i, &vid) in self.outputs.iter().enumerate() {
            if i > 0 {
                s.push_str(", ");
            }
            s.push_str(&format!("{}", self.values[vid as usize].name));
        }
        s.push_str("]\n");
        let topo = self.topo_order();
        for (idx, nid) in topo.iter().enumerate() {
            let n = &self.nodes[*nid as usize];
            let dom = if n.domain.is_empty() {
                "ai.onnx"
            } else {
                &n.domain
            };
            let in_names: alloc::vec::Vec<_> = n
                .inputs
                .iter()
                .map(|&v| self.values[v as usize].name.as_str())
                .collect();
            let out_names: alloc::vec::Vec<_> = n
                .outputs
                .iter()
                .map(|&v| self.values[v as usize].name.as_str())
                .collect();

            s.push_str(&format!(
                "[{}] {}::{}  {} -> {}\n",
                idx,
                dom,
                n.op_type,
                fmt_list(&in_names),
                fmt_list(&out_names),
            ));

            if opts.show_attrs && !n.attrs.is_empty() {
                s.push_str("     attrs: ");
                for (i, (k, v)) in n.attrs.iter().enumerate() {
                    if i > 0 {
                        s.push_str(", ");
                    }
                    s.push_str(&format!("{}={}", k, fmt_attr(v, opts.const_max_elems)));
                }
                s.push('\n');
            }
        }

        if opts.show_constants {
            s.push_str("constants:\n");
            for v in &self.values {
                if let Some(t) = &v.constant {
                    s.push_str(&format!(
                        "  {}: {}{}\n",
                        v.name,
                        fmt_ty(&t.dtype, &t.shape),
                        fmt_const_preview(t, opts.const_max_elems),
                    ));
                }
            }
        }

        s
    }

    pub fn to_mermaid(&self) -> alloc::string::String {
        use alloc::format;
        let mut m = alloc::string::String::from("flowchart LR\n");

        for v in &self.values {
            let label = if let Some(t) = &v.constant {
                format!("{}\\n{}", v.name, fmt_ty(&t.dtype, &t.shape))
            } else {
                v.name.clone()
            };
            m.push_str(&format!("  V{}[\"{}\"]\n", v.id, escape_mermaid(&label)));
        }

        for n in &self.nodes {
            let dom = if n.domain.is_empty() {
                "ai.onnx"
            } else {
                &n.domain
            };
            m.push_str(&format!("  N{}((\"{}::{}\"))\n", n.id, dom, n.op_type));
            for &iv in &n.inputs {
                m.push_str(&format!("  V{} --> N{}\n", iv, n.id));
            }
            for &ov in &n.outputs {
                m.push_str(&format!("  N{} --> V{}\n", n.id, ov));
            }
        }
        m
    }
}

impl fmt::Display for super::graph::Graph {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let opts = PrettyOpts {
            show_attrs: true,
            show_constants: false,
            const_max_elems: 8,
        };
        write!(f, "{}", self.pretty(opts))
    }
}

fn fmt_list<'a>(xs: &[&'a str]) -> alloc::string::String {
    use alloc::string::ToString;
    let mut s = "[".to_string();
    for (i, x) in xs.iter().enumerate() {
        if i > 0 {
            s.push_str(", ");
        }
        s.push_str(x);
    }
    s.push(']');
    s
}

fn fmt_ty(dt: &crate::ir::tensor::DType, shape: &[usize]) -> alloc::string::String {
    use alloc::string::ToString;
    let mut s = match dt {
        crate::ir::tensor::DType::F16 => "f16".to_string(),
        crate::ir::tensor::DType::F32 => "f32".to_string(),
        crate::ir::tensor::DType::F64 => "f64".to_string(),
        crate::ir::tensor::DType::I8 => "i8".to_string(),
        crate::ir::tensor::DType::I16 => "i16".to_string(),
        crate::ir::tensor::DType::I32 => "i32".to_string(),
        crate::ir::tensor::DType::I64 => "i64".to_string(),
        crate::ir::tensor::DType::U8 => "u8".to_string(),
        crate::ir::tensor::DType::U16 => "u16".to_string(),
        crate::ir::tensor::DType::U32 => "u32".to_string(),
        crate::ir::tensor::DType::U64 => "u64".to_string(),
        crate::ir::tensor::DType::Bool => "bool".to_string(),
        crate::ir::tensor::DType::BF16 => "bf16".to_string(),
    };
    s.push('[');
    for (i, d) in shape.iter().enumerate() {
        if i > 0 {
            s.push(',');
        }
        s.push_str(&d.to_string());
    }
    s.push(']');
    s
}

fn fmt_attr(a: &crate::ir::attribute::AttrValue, max: usize) -> alloc::string::String {
    use alloc::format;
    match a {
        crate::ir::attribute::AttrValue::Int(i) => format!("{}", i),
        crate::ir::attribute::AttrValue::Float(x) => format!("{}", x),
        crate::ir::attribute::AttrValue::Ints(v) => fmt_vec_i64(v, max),
        crate::ir::attribute::AttrValue::Floats(v) => fmt_vec_f32(v, max),
        crate::ir::attribute::AttrValue::Bytes(b) => format!("bytes(len={})", b.len()),
        crate::ir::attribute::AttrValue::Tensor(t) => {
            format!("Tensor{}", fmt_ty(&t.dtype, &t.shape))
        }
    }
}

fn fmt_vec_i64(v: &alloc::vec::Vec<i64>, max: usize) -> alloc::string::String {
    use alloc::string::ToString;
    let mut s = "[".to_string();
    for (i, x) in v.iter().take(max).enumerate() {
        if i > 0 {
            s.push_str(", ");
        }
        s.push_str(&x.to_string());
    }
    if v.len() > max {
        s.push_str(", ...");
    }
    s.push(']');
    s
}

fn fmt_vec_f32(v: &alloc::vec::Vec<f32>, max: usize) -> alloc::string::String {
    use alloc::string::ToString;
    let mut s = "[".to_string();
    for (i, x) in v.iter().take(max).enumerate() {
        if i > 0 {
            s.push_str(", ");
        }
        s.push_str(&x.to_string());
    }
    if v.len() > max {
        s.push_str(", ...");
    }
    s.push(']');
    s
}

fn fmt_const_preview(t: &crate::ir::tensor::Tensor, max: usize) -> alloc::string::String {
    use alloc::format;
    if matches!(t.dtype, crate::ir::tensor::DType::F32) {
        let xs = t.as_slice::<f32>();
        let shown = core::cmp::min(xs.len(), max);
        let mut s = alloc::string::String::from(" = [");
        for i in 0..shown {
            if i > 0 {
                s.push_str(", ");
            }
            s.push_str(&format!("{}", xs[i]));
        }
        if xs.len() > shown {
            s.push_str(", ...");
        }
        s.push(']');
        s
    } else {
        alloc::format!(" ({} bytes)", t.data.len())
    }
}

fn escape_mermaid(s: &str) -> alloc::string::String {
    s.replace('"', "\\\"")
}
